# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/22 10:46 
@Author  : ZhangShenao 
@File    : 1.使用自定义向量数据库.py 
@Desc    : 使用自定义的InMemoryVectorStore
"""

import os

import dotenv
from langchain_openai import OpenAIEmbeddings

from in_memory_vector_store import InMemoryVectorStore

# 加载环境变量
dotenv.load_dotenv()

# 创建Embeddings
embeddings = OpenAIEmbeddings(
    model='text-embedding-3-small',  # 选择text-embedding-3-small模型
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 创建Pinecone向量数据库
vector_store = InMemoryVectorStore(
    embedding=embeddings,  # 指定Embedding模型
)

texts = [
    '笨笨是一只很喜欢睡觉的猫咪',
    '我喜欢在夜晚听音乐，这让我感到放松。',
    '猫咪在窗台上打盹，看起来非常可爱。',
    '学习新技能是每个人都应该追求的目标。',
    '我最喜欢的食物是意大利面，尤其是番茄酱的那种。',
    '昨晚我做了一个奇怪的梦，梦见自己在太空飞行。',
    '我的手机突然关机了，让我有些焦虑。',
    '阅读是我每天都会做的事情，我觉得很充实。',
    '他们一起计划了一次周末的野餐，希望天气能好。',
    '我的狗喜欢追逐球，看起来非常开心。',
]

meta_datas = [
    {'page': 1},
    {'page': 2},
    {'page': 3},
    {'page': 4},
    {'page': 5},
    {'page': 6},
    {'page': 7},
    {'page': 8},
    {'page': 9},
    {'page': 10},
]

# 向自定义向量数据库中插入数据
ids = vector_store.add_texts(
    texts=texts,  # 指定文本内容
    metadatas=meta_datas,  # 指定元数据
)
print(ids)

# 进行相似性检索
search_result = vector_store.similarity_search(
    query='我养了一只猫叫笨笨',  # 指定查询文本
    k=2,  # 指定返回的文档数量
)
print(f'search_result: {search_result}')
